//Patric Zhao:  patric.zhao@gmail.com

#include <chrono>
#include <iostream>
#include <CL/sycl.hpp>

#define random_float() (rand() / double(RAND_MAX))
#define random_int(n) (rand() % n)

#define tileN 16
#define sizeM 32

using namespace std;
using namespace sycl;

double cpu_kernel(float* x, int* mask, float* weight, float* loss, const int K, const int M, const int N) {

    double duration = 0.0;
    std::chrono::high_resolution_clock::time_point s, e;

    // Single Thread Computation in CPU
    s = std::chrono::high_resolution_clock::now();
    float sum = 0, max = 0, y = 0;
    for (int i = 0; i < K * N; i++) {
        sum = 0, max = 0, y = 0;
        for (int m = 0; m < M; m++)
            if (x[i * M + m] > max)
                max = x[i * M + m];

        for (int m = 0; m < M; m++)
            sum += exp(x[i * M + m] - max);

        y = x[i * M + mask[i]] - max - log(sum);

        loss[i] = -y * weight[i];
    }
    e = std::chrono::high_resolution_clock::now();
    duration = std::chrono::duration<float, std::milli>(e - s).count();

    return(duration);
}

// return execution time
double gpu_kernel(float* x, int* mask, float* weight, float* loss, const int K, const int M, const int N, sycl::queue& q) {

    // define the workgroup size and mapping
    auto grid_cols = (K * N + tileN - 1) / tileN * tileN;
    auto local_ndrange = range<1>(tileN);
    auto global_ndrange = range<1>(grid_cols);

    double duration = 0.0f;
    auto e = q.submit([&](sycl::handler& h) {
        h.parallel_for<class k_name_t>(
            sycl::nd_range<1>(global_ndrange, local_ndrange), [=](sycl::nd_item<1> index) {
                int i = index.get_global_id();
                float sum = 0, max = 0, y = 0;

                // cache x
                float x_local[sizeM] = {};

                for (int m = 0; m < M; m++) {
                    x_local[m] = x[i * M + m];
                    if (x_local[m] > max)
                        max = x_local[m];
                }

                for (int m = 0; m < M; m++)
                    sum += exp(x_local[m] - max);

                y = x_local[mask[i]] - max - log(sum);

                loss[i] = -y * weight[i];
            });
        });
    e.wait();

    duration += (e.get_profiling_info<info::event_profiling::command_end>() -
        e.get_profiling_info<info::event_profiling::command_start>()) / 1000.0f / 1000.0f;

    return(duration);
}

int verify(float* cpu_res, float* gpu_res, int length) {
    int err = 0;
    for (int i = 0; i < length; i++) {
        if (fabs(cpu_res[i] - gpu_res[i]) > 1e-3) {
            err++;
            printf("\n%lf, %lf", cpu_res[i], gpu_res[i]);
        }
    }
    return(err);
}

int cross_entropy(const int K, const int M, const int N, const int iterations, sycl::queue& q) {
    auto x = malloc_shared<float>(K * M * N, q);
    auto mask = malloc_shared<int>(K * N, q);
    auto weight = malloc_shared<float>(K * N, q);
    for (int i = 0; i < K * M * N; i++)
        x[i] = 10 * random_float();
    for (int i = 0; i < K * N; i++) {
        mask[i] = random_int(M);
        weight[i] = random_float();
    }

    auto loss_device = malloc_shared<float>(K * N, q);
    auto loss_host = malloc_host<float>(K * N, q);

    double gpu_duration = 0.0f;
    double cpu_duration = 0.0f;

    // GPU compuation and timer 
    int warmup = 10;
    for (int run = 0; run < iterations + warmup; run++) {
        float duration = gpu_kernel(x, mask, weight, loss_device, K, M, N, q);
        if (run >= warmup) gpu_duration += duration;
    }
    gpu_duration = gpu_duration / iterations;

    // CPU compuation and timer 
    warmup = 2;
    for (int run = 0; run < iterations / 2 + warmup; run++) {
        float duration = cpu_kernel(x, mask, weight, loss_host, K, M, N);
        if (run >= warmup) cpu_duration += duration;
    }
    cpu_duration = cpu_duration / (iterations / 2);

    // Compare the resutls of CPU and GPU 
    int errCode = 0;
    errCode = verify(loss_host, loss_device, K * N);
    if (errCode > 0) printf("\nThere are %d errors\n", errCode);

    printf("for tileN = %d: \n"
        "GPU Computation Time = %lf (ms); \n"
        "CPU Computaiton Time = %lf (ms); \n",
        tileN, gpu_duration, cpu_duration);

    free(x, q);
    free(mask, q);
    free(weight, q);
    free(loss_host, q);
    free(loss_device, q);

    return(errCode);
}

int main() {

    auto propList = cl::sycl::property_list{ cl::sycl::property::queue::enable_profiling() };
    queue my_gpu_queue(cl::sycl::gpu_selector{}, propList);

    //   int errCode = gemm(1024, 1024, 1024, 4, 10, my_gpu_queue);
    int errCode = cross_entropy(128, 32, 8192, 10, my_gpu_queue);
    cout << "success" << std::endl;
    return(errCode);
}
